% Compute the global solution using a paramerized expectations algorithm
% Dynamic bank run model
% Bank leverage choice only
% TFP shock

clear
close all
warning off

%% Generate random shocks

t_sim  = 10000;                   % simulation periods
rng(1)                            % same random shock 
shocks = normrnd(0,1,t_sim+1,1);  % random shock
%shocks = dlmread(fullfile('from_matlab_to_fortran/shocks.txt'));
t_drop = 500;                     % number of initial periods to drop in regression

%% Read parameters and steady state

filename = '../linear_model/model_results.mat';
load(filename);
Np = M_.param_nbr;                                            % Number of deep parameters.
for i = 1:Np                                                  % Loop...
  paramname = deblank(M_.param_names(i,:));                   % Get the name of parameter i. 
  eval([ paramname ' = M_.params(' int2str(i) ');']);         % Get the value of parameter i.
end

% 'hU_SS  ' SS labor
% 'LU_SS  ' SS leverage
% 'PU_SS  ' SS run probability
beta    = beta_p;% 'beta_p ' beta
nu      = nu_p;% 'nu_p   ' nu
alpha   = alpha_p;% 'alpha_p' alpha
delta   = delta_p;% 'delta_p' delta
lambda  = lam_p;% 'lam_p  ' lambda
chi0    = chi0_p;% 'chi0_p ' chi0
chi1    = chi1_p;% 'chi1_p ' chi1
rho_a   = rhoa_p;% 'rhoa_p ' rho_a
sigma_a = siga_p;% 'siga_p ' sigma_a

y_ss  = oo_.mean(1);
i_ss  = oo_.mean(2);
xi_ss = oo_.mean(3);
pr_ss = oo_.mean(4);
psi_p = oo_.mean(5);
n0_p  = oo_.mean(6);
gam_p = oo_.mean(7);
k_ss  = oo_.mean(8);
n_ss  = oo_.mean(9);
l_ss  = oo_.mean(10);
rb_ss = oo_.mean(11);
rkhats_ss = oo_.mean(12);
a_ss  = oo_.mean(13);
c_ss  = oo_.mean(14);
h_ss  = oo_.mean(15);

psi   = psi_p;
n0    = n0_p;
gamma = gam_p;

sigma_rk = (1+nu_p)/(1+alpha_p*nu_p)*siga_p;

keep   = lambda*(1-gamma);
lambda = 0.15/0.85;
gamma  = 1-keep/lambda;

lam_p  = lambda;
gam_p  = gamma;

%% Prepare other parameters and polynomial matrix

npf  = 7;  % number of policy functions to approximate
nstv = 4;  % number of state variables
degp = 3;  % degree of polynomials

expmx         = fn_expmx(nstv,degp);    % exponent matrix
expmx_fortran = expmx';

npol = size(expmx,1);

%% initial guess for etas

init_eta1_rhsee = zeros(npol,1);
init_eta2_rhsee = zeros(npol,1);
init_eta3_rhsee = zeros(npol,1);
init_eta4_rhsee = zeros(npol,1);
%eta1_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta1_rhsee.txt'));
%eta2_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta2_rhsee.txt'));
%eta3_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta3_rhsee.txt'));
%eta4_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta4_rhsee.txt'));
%eta1_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta1_rhsee.txt'));
%eta2_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta2_rhsee.txt'));
%eta3_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta3_rhsee.txt'));
%eta4_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta4_rhsee.txt'));
eta1_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta1_rhsee.txt'));
eta2_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta2_rhsee.txt'));
eta3_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta3_rhsee.txt'));
eta4_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta4_rhsee.txt'));
for i=1:npol
    init_eta1_rhsee(i) = eta1_import(i);
    init_eta2_rhsee(i) = eta2_import(i);
    init_eta3_rhsee(i) = eta3_import(i);
    init_eta4_rhsee(i) = eta4_import(i);
end
%init_eta4_rhsee = init_eta1_rhsee;

init_eta1_rbar = zeros(npol,1);
init_eta2_rbar = zeros(npol,1);
init_eta3_rbar = zeros(npol,1);
%eta1_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta1_rbar.txt'));
%eta2_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta2_rbar.txt'));
%eta3_import = dlmread(fullfile('from_matlab_to_fortran/solved_eta3_rbar.txt'));
%eta1_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta1_rbar.txt'));
%eta2_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta2_rbar.txt'));
%eta3_import = dlmread(fullfile('from_fortran_to_matlab/temp_eta3_rbar.txt'));
eta1_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta1_rbar.txt'));
eta2_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta2_rbar.txt'));
eta3_import = dlmread(fullfile('from_fortran_to_matlab/solved_eta3_rbar.txt'));
for i=1:npol
    init_eta1_rbar(i) = eta1_import(i);
    init_eta2_rbar(i) = eta2_import(i);
    init_eta3_rbar(i) = eta3_import(i);
end

%% Set parameters for numerical integration and iterations

nd = 70;  % nd percent drop in net worth in a crisis

n_trapz = 301;
z_max   = 5;
z_min   = -z_max;
z_trapz = linspace(z_min,z_max,n_trapz);
z_trapz = z_trapz';

%% Prepare for welfare analysis

length = 3000;
repeat = 1000;
welfare_integer = [length,repeat];
save('from_matlab_to_fortran/welfare_integer.txt','welfare_integer','-ascii','-double');

%filename = '../no_policy_3rd/no_policy.mat';
filename = matfile('../no_policy_3rd/no_policy.mat');
stss_n   = filename.stss_n;
stss_k   = filename.stss_k;
stss_rb  = filename.stss_rb;

n1   = stss_n;
k1   = stss_k;
rb1  = stss_rb;
welfare_states = [n1,k1,rb1];
save('from_matlab_to_fortran/welfare_states.txt','welfare_states','-ascii','-double');

pn       = 0.20;
thre_l   = stss_k/stss_n * 1.15;
thre_n   = stss_n * 1;
thre_pib = 0;
policy   = [pn,thre_l,thre_n,thre_pib];
save('from_matlab_to_fortran/policy_parameters.txt','policy','-ascii','-double');

iter_max = 100;
tol      = 1e-3;   % convergence criterion
adj      = 0.5;   % adjustment speed of polynomial coefficients, the higher the larger update and the faster.

%% Prepare for fortran

parameters    = [hU_SS,LU_SS,PU_SS,beta_p,nu_p,alpha_p,delta_p,lam_p,chi0_p,chi1_p,rhoa_p,siga_p,nd];
save('from_matlab_to_fortran/parameters.txt','parameters','-ascii');

parameters_ss = [c_ss,y_ss,i_ss,k_ss,h_ss,n_ss,l_ss,rb_ss,rkhats_ss,a_ss,xi_ss,pr_ss,psi_p,n0_p,gam_p,sigma_rk];
save('from_matlab_to_fortran/ss_values.txt','parameters_ss','-ascii');

parameters_integer = [npf,nstv,degp,npol,iter_max,n_trapz,t_sim,t_drop];
save('from_matlab_to_fortran/parameters_integer.txt','parameters_integer','-ascii');

parameters_solution = [tol,adj];
save('from_matlab_to_fortran/parameters_solution.txt','parameters_solution','-ascii');

save('from_matlab_to_fortran/polynomial_matrix.txt','expmx_fortran','-ascii','-double');
save('from_matlab_to_fortran/z_trapz.txt','z_trapz','-ascii','-double');
save('from_matlab_to_fortran/shocks.txt','shocks','-ascii','-double');

%% run fortran

notsolved = zeros(1,1);

welfare = zeros(1,1);

% reset initial guess when nss is initialized at 1
eta1_rhsee = init_eta1_rhsee;
eta2_rhsee = init_eta2_rhsee;
eta3_rhsee = init_eta3_rhsee;
eta4_rhsee = init_eta4_rhsee;
eta1_rbar  = init_eta1_rbar;
eta2_rbar  = init_eta2_rbar;
eta3_rbar  = init_eta3_rbar;

for nss=1:1
        
    eta1_rhsee = init_eta1_rhsee;
    eta2_rhsee = init_eta2_rhsee;
    eta3_rhsee = init_eta3_rhsee;
    eta4_rhsee = init_eta4_rhsee;
    eta1_rbar  = init_eta1_rbar;
    eta2_rbar  = init_eta2_rbar;
    eta3_rbar  = init_eta3_rbar;
    
    for lss=11:11
    
        pn       = 0.0;
        thre_pib = 0;

%        thre_l  = k_stss_nopol/n_stss_nopol * ( 1.59 - 0.01*(lss-1) );
        thre_l  = 16 - 0.1*(lss-1);
        thre_n  = stss_n * ( 0 );
        X = [ '############################ lss=',num2str(thre_l),' ############################' ];
        disp(X)

        policy   = [pn,thre_l,thre_n,thre_pib];
        save('from_matlab_to_fortran/policy_parameters.txt','policy','-ascii','-double');

for i=1:iter_max
    
    notsolved(lss,nss) = 0;
    
    % export etas to fortran
    save('from_matlab_to_fortran/eta1_rhsee.txt','eta1_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta2_rhsee.txt','eta2_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta3_rhsee.txt','eta3_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta4_rhsee.txt','eta4_rhsee','-ascii','-double');
    save('from_matlab_to_fortran/eta1_rbar.txt','eta1_rbar','-ascii','-double');
    save('from_matlab_to_fortran/eta2_rbar.txt','eta2_rbar','-ascii','-double');
    save('from_matlab_to_fortran/eta3_rbar.txt','eta3_rbar','-ascii','-double');
    
    system('x64\Release\f90_bankrun_pea.exe');
        
    numbers    = dlmread(fullfile('from_fortran_to_matlab/numbers.txt'));
    ee_error   = dlmread(fullfile('from_fortran_to_matlab/ee_error.txt'));
    ee_error_info = dlmread(fullfile('from_fortran_to_matlab/ee_error_info.txt'));
    rbv_maxmin = dlmread(fullfile('from_fortran_to_matlab/rbv.txt'));

    num_norun = numbers(1);
    num_run   = numbers(2);
    num_nega  = numbers(3);
    num_pol   = numbers(4);
    
%    system('x64/Debug/calibration2.exe');

    % import rhsee and rbar derived, and sequence of state variables x
    f = dir('from_fortran_to_matlab/rhsee_norun.txt');
    if f.bytes>10
        rhsee_norun = dlmread(fullfile('from_fortran_to_matlab/rhsee_norun.txt'));
        rbar_norun  = dlmread(fullfile('from_fortran_to_matlab/rbar_norun.txt'));
        x_norun     = dlmread(fullfile('from_fortran_to_matlab/x_norun.txt'));
        num_norun = size(x_norun,1)/npol;
        x_norun   = reshape(x_norun,num_norun,npol);
        eta1_rhsee_new = (x_norun'*x_norun)\x_norun'*rhsee_norun;
        eta1_rbar_new  = (x_norun'*x_norun)\x_norun'*rbar_norun;
        con_lev(1) = max( abs( (exp(x_norun*eta1_rhsee_new) - exp(x_norun*eta1_rhsee)) ./ exp(x_norun*eta1_rhsee) ) );
        con_lev(5) = max( abs( (exp(x_norun*eta1_rbar_new ) - exp(x_norun*eta1_rbar )) ./ exp(x_norun*eta1_rbar ) ) );
        rsq(1) = 1 - sum((rhsee_norun - x_norun*eta1_rhsee).^2) ./ sum((rhsee_norun-mean(rhsee_norun)).^2);
        rsq(5) = 1 - sum((rbar_norun  - x_norun*eta1_rbar ).^2) ./ sum((rbar_norun -mean(rbar_norun )).^2);
    else
        eta1_rhsee_new = eta1_rhsee;
        eta1_rbar_new  = eta1_rbar;
        con_lev(1) = 0;
        con_lev(5) = 0;
        rsq(1) = 0;
        rsq(5) = 0;
    end

    f = dir('from_fortran_to_matlab/rhsee_run.txt');
    if f.bytes>10
        rhsee_run = dlmread(fullfile('from_fortran_to_matlab/rhsee_run.txt'));
        rbar_run  = dlmread(fullfile('from_fortran_to_matlab/rbar_run.txt'));
        x_run     = dlmread(fullfile('from_fortran_to_matlab/x_run.txt'));
        num_run = size(x_run,1)/npol;
        x_run   = reshape(x_run,num_run,npol);
        eta2_rhsee_new = (x_run'*x_run)\x_run'*rhsee_run;
        eta2_rbar_new  = (x_run'*x_run)\x_run'*rbar_run;
        con_lev(2) = max( abs( (exp(x_run  *eta2_rhsee_new) - exp(x_run  *eta2_rhsee)) ./ exp(x_run  *eta2_rhsee) ) );
        con_lev(6) = max( abs( (exp(x_run  *eta2_rbar_new ) - exp(x_run  *eta2_rbar )) ./ exp(x_run  *eta2_rbar ) ) );
        rsq(2) = 1 - sum((rhsee_run   - x_run  *eta2_rhsee).^2) ./ sum((rhsee_run  -mean(rhsee_run  )).^2);
        rsq(6) = 1 - sum((rbar_run    - x_run  *eta2_rbar ).^2) ./ sum((rbar_run   -mean(rbar_run   )).^2);
    else
        eta2_rhsee_new = eta2_rhsee;
        eta2_rbar_new  = eta2_rbar;
        con_lev(2) = 0;
        con_lev(6) = 0;
        rsq(2) = 0;
        rsq(6) = 0;
    end

    f = dir('from_fortran_to_matlab/rhsee_nega.txt');
    if f.bytes>10
        rhsee_nega = dlmread(fullfile('from_fortran_to_matlab/rhsee_nega.txt'));
        rbar_nega  = dlmread(fullfile('from_fortran_to_matlab/rbar_nega.txt'));
        x_nega     = dlmread(fullfile('from_fortran_to_matlab/x_nega.txt'));
        num_nega = size(x_nega,1)/npol;
        x_nega   = reshape(x_nega,num_nega,npol);
        eta3_rhsee_new = (x_nega'*x_nega)\x_nega'*rhsee_nega;
        eta3_rbar_new  = (x_nega'*x_nega)\x_nega'*rbar_nega;
        con_lev(3) = max( abs( (exp(x_nega *eta3_rhsee_new) - exp(x_nega *eta3_rhsee)) ./ exp(x_nega *eta3_rhsee) ) );
        con_lev(7) = max( abs( (exp(x_nega *eta3_rbar_new ) - exp(x_nega *eta3_rbar )) ./ exp(x_nega *eta3_rbar) ) );
        rsq(3) = 1 - sum((rhsee_nega  - x_nega *eta3_rhsee).^2) ./ sum((rhsee_nega -mean(rhsee_nega )).^2);
        rsq(7) = 1 - sum((rbar_nega   - x_nega *eta3_rbar ).^2) ./ sum((rbar_nega  -mean(rbar_nega  )).^2);
    else
        eta3_rhsee_new = eta3_rhsee;
        eta3_rbar_new  = eta3_rbar;
        con_lev(3) = 0;
        con_lev(7) = 0;
        rsq(3) = 0;
        rsq(7) = 0;
    end

    f = dir('from_fortran_to_matlab/rhsee_pol.txt');
    if f.bytes>10
        rhsee_pol = dlmread(fullfile('from_fortran_to_matlab/rhsee_pol.txt'));
        x_pol     = dlmread(fullfile('from_fortran_to_matlab/x_pol.txt'));
        num_pol = size(x_pol,1)/npol;
        x_pol   = reshape(x_pol,num_pol,npol);
        eta4_rhsee_new  = (x_pol'*x_pol)\x_pol'*rhsee_pol;
        con_lev(4) = max( abs( (exp(x_pol  *eta4_rhsee_new ) - exp(x_pol  *eta4_rhsee )) ./ exp(x_pol  *eta4_rhsee ) ) );
        rsq(4)     = 1 - sum((rhsee_pol    - x_pol  *eta4_rhsee ).^2) ./ sum((rhsee_pol   -mean(rhsee_pol   )).^2);
    else
        eta4_rhsee_new  = eta4_rhsee;
        con_lev(4) = 0;
        rsq(4)     = 0;
    end
        
    % check convergence
    
    con_lev = zeros(7,1);
    rsq     = zeros(7,1);
    
    con_lev(1) = mean( abs( (exp(x_norun*eta1_rhsee_new) - exp(x_norun*eta1_rhsee)) ./ exp(x_norun*eta1_rhsee) ) );
    con_lev(2) = mean( abs( (exp(x_run  *eta2_rhsee_new) - exp(x_run  *eta2_rhsee)) ./ exp(x_run  *eta2_rhsee) ) );
    con_lev(3) = mean( abs( (exp(x_nega *eta3_rhsee_new) - exp(x_nega *eta3_rhsee)) ./ exp(x_nega *eta3_rhsee) ) );
    if num_pol > 0
        con_lev(4) = max( abs( (exp(x_pol  *eta4_rhsee_new) - exp(x_pol  *eta4_rhsee)) ./ exp(x_pol  *eta4_rhsee) ) );
    end
    con_lev(5) = mean( abs( (exp(x_norun*eta1_rbar_new ) - exp(x_norun*eta1_rbar )) ./ exp(x_norun*eta1_rbar ) ) );
    con_lev(6) = mean( abs( (exp(x_run  *eta2_rbar_new ) - exp(x_run  *eta2_rbar )) ./ exp(x_run  *eta2_rbar ) ) );
    con_lev(7) = mean( abs( (exp(x_nega *eta3_rbar_new ) - exp(x_nega *eta3_rbar )) ./ exp(x_nega *eta3_rbar) ) );

    rsq(1) = 1 - sum((rhsee_norun - x_norun*eta1_rhsee_new).^2) ./ sum((rhsee_norun-mean(rhsee_norun)).^2);
    rsq(2) = 1 - sum((rhsee_run   - x_run  *eta2_rhsee_new).^2) ./ sum((rhsee_run  -mean(rhsee_run  )).^2);
    rsq(3) = 1 - sum((rhsee_nega  - x_nega *eta3_rhsee_new).^2) ./ sum((rhsee_nega -mean(rhsee_nega )).^2);
    if num_pol > 0
        rsq(4) = 1 - sum((rhsee_pol   - x_pol  *eta4_rhsee_new).^2) ./ sum((rhsee_pol  -mean(rhsee_pol  )).^2);
    end
    rsq(5) = 1 - sum((rbar_norun  - x_norun*eta1_rbar_new ).^2) ./ sum((rbar_norun -mean(rbar_norun )).^2);
    rsq(6) = 1 - sum((rbar_run    - x_run  *eta2_rbar_new ).^2) ./ sum((rbar_run   -mean(rbar_run   )).^2);
    rsq(7) = 1 - sum((rbar_nega   - x_nega *eta3_rbar_new ).^2) ./ sum((rbar_nega  -mean(rbar_nega  )).^2);

    % display results
    fprintf('#################### iteration:')
    disp(i)

    [max_con,max_con_ind] = max(con_lev(1:4));
    X = ['max gap in convergence: ',num2str(max_con),', at ',sprintf('%d',round(max_con_ind))];    
    disp(X)
    
    disp('R-square:')
    X = ['RHSEE: ',num2str(rsq(1:4)')];
    disp(X)
    X = ['Number of periods:  ',num2str(numbers)];
    disp(X)
    X = ['EE error (mean,max) ',num2str(mean(ee_error)),' ',num2str(max(ee_error))];
    disp(X)
    
    if max(con_lev(1:4)) < 0.01 && max(ee_error)<-2
        break
    end
    
    % update etas
    eta1_rhsee = adj*eta1_rhsee_new + (1-adj)*eta1_rhsee;
    eta2_rhsee = adj*eta2_rhsee_new + (1-adj)*eta2_rhsee;
    eta3_rhsee = adj*eta3_rhsee_new + (1-adj)*eta3_rhsee;
    eta4_rhsee = adj*eta4_rhsee_new + (1-adj)*eta4_rhsee;
    eta1_rbar  = adj*eta1_rbar_new  + (1-adj)*eta1_rbar;
    eta2_rbar  = adj*eta2_rbar_new  + (1-adj)*eta2_rbar;
    eta3_rbar  = adj*eta3_rbar_new  + (1-adj)*eta3_rbar;

end

if notsolved(lss,nss) == 1
    break
end

% import from file 'temp' as variables 'solved_etas'
solved_eta1_rhsee = dlmread(fullfile('from_fortran_to_matlab/temp_eta1_rhsee.txt'));
solved_eta2_rhsee = dlmread(fullfile('from_fortran_to_matlab/temp_eta2_rhsee.txt'));
solved_eta3_rhsee = dlmread(fullfile('from_fortran_to_matlab/temp_eta3_rhsee.txt'));
solved_eta4_rhsee = dlmread(fullfile('from_fortran_to_matlab/temp_eta4_rhsee.txt'));
solved_eta1_rbar  = dlmread(fullfile('from_fortran_to_matlab/temp_eta1_rbar.txt'));
solved_eta2_rbar  = dlmread(fullfile('from_fortran_to_matlab/temp_eta2_rbar.txt'));
solved_eta3_rbar  = dlmread(fullfile('from_fortran_to_matlab/temp_eta3_rbar.txt'));

% save variables 'solved_etas' as file 'solved'
save('from_fortran_to_matlab/solved_eta1_rhsee.txt','solved_eta1_rhsee','-ascii','-double');
save('from_fortran_to_matlab/solved_eta2_rhsee.txt','solved_eta2_rhsee','-ascii','-double');
save('from_fortran_to_matlab/solved_eta3_rhsee.txt','solved_eta3_rhsee','-ascii','-double');
save('from_fortran_to_matlab/solved_eta4_rhsee.txt','solved_eta4_rhsee','-ascii','-double');
save('from_fortran_to_matlab/solved_eta1_rbar.txt' ,'solved_eta1_rbar' ,'-ascii','-double');
save('from_fortran_to_matlab/solved_eta2_rbar.txt' ,'solved_eta2_rbar' ,'-ascii','-double');
save('from_fortran_to_matlab/solved_eta3_rbar.txt' ,'solved_eta3_rbar' ,'-ascii','-double');

command='./main_welfare.exe';
[status,cmdout] = unix(command);
welfare(lss,nss) = dlmread(fullfile('from_fortran_to_matlab/utility_compare.txt'))

%if rem(lss,10)==0 && rem(nss,10)==0

%    filename = ['policy_threL_',num2str(160 - 1*(lss-1))];
%    save(filename)

%end

    end

end

return
